# -*- coding: utf-8 -*-

import cv2
import os
import sys

from PIL import Image

from photoClassify.project.faceClasiify.config import peoples_dir_path, model_name, video_path, image_size


def CatchPICFromVideo(name):
    while True:
        cv2.namedWindow('录入')

        # 视频来源
        print('---数据保存路径:'+peoples_dir_path+'\n')
        source = input('输入要截取的视频(文件名+格式)或摄像头(0)')

        # 捕获指定摄像头或视频资源
        if source == '0':
            cap = cv2.VideoCapture(0)
        elif source.__len__() > 0:
            cap = cv2.VideoCapture(video_path + source)
        else:
            continue # 跳过当前循环块中的剩余语句，然后继续进行下一轮循环

        # OpenCV分类器
        classfier = cv2.CascadeClassifier(
            "/Users/mc/PycharmProjects/TF-vision/venv/lib/python3.7/site-packages/cv2/data/haarcascade_frontalface_alt2.xml")

        num = 0
        while cap.isOpened():
            ok, frame = cap.read()  # 读取一帧数据
            if not ok:
                break

            grey = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)  # 将当前桢图像转换成灰度图像

            # 检测
            faceRects = classfier.detectMultiScale(grey, scaleFactor=1.2, minNeighbors=3, minSize=(40, 40),
                                                   maxSize=(96, 96))
            if len(faceRects) > 0:  # 大于0则检测到人脸
                for faceRect in faceRects:  # 单独框出每一张人脸
                    x, y, w, h = faceRect

                    # # 将当前帧保存为图片
                    img_name = '%s/%d.jpg' % (peoples_dir_path, num)
                    image = frame[y: y + h, x: x + w]
                    cv2.imwrite(img_name, image)
                    num += 1
                    # 画出矩形框
                    cv2.rectangle(frame, (x, y), (x + w, y + h), (101, 101, 255), 2)
                    # cv2.putText(frame, 'num:%d' % (num), (x + 30, y + 30), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 4)

            # 显示图像
            cv2.imshow('录入', frame)
            c = cv2.waitKey(30)
            if c & 0xFF == ord('q'):
                break

                # 释放摄像头并销毁所有窗口
        cap.release()
        cv2.destroyAllWindows()


# 截取视频中人脸保存到文件夹下
# 每个模块都有一个__name__属性，当其值是'__main__'时，表明该模块自身在运行，否则是被引入
if __name__ == '__main__':

    if len(sys.argv) != 2:
        CatchPICFromVideo("none")
    else:
        CatchPICFromVideo(sys.argv[1])
